from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from jose import JWTError, jwt
from app.core.config import settings
from app.core.security import ALGORITHM
from app.database.database import SessionLocal
from starlette.responses import JSONResponse
import logging

# 创建日志记录器
logger = logging.getLogger(__name__)

class WorkAreaMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # 排除不需要工作区验证的路径
        if request.url.path in ["/docs", "/redoc", "/openapi.json"] or request.url.path.startswith("/api/auth"):
            return await call_next(request)
            
        try:
            # 获取 token
            auth_header = request.headers.get("Authorization")
            if not auth_header or not auth_header.startswith("Bearer "):
                return await call_next(request)

            token = auth_header.split(" ")[1]
            
            # 验证 token
            try:
                # 使用 settings.SECRET_KEY
                payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[ALGORITHM])
                user_id: str = payload.get("sub")
                if not user_id:
                    return await call_next(request)
                
                # 添加日志
                logger.info(f"Token payload: {payload}")
                logger.info(f"User ID from token: {user_id}")
                
                # 获取用户信息
                db = SessionLocal()
                try:
                    from app.database.models.base import User
                    user = db.query(User).filter(User.id == int(user_id)).first()
                    
                    # 添加日志
                    logger.info(f"Found user: {user.username if user else None}")
                    
                    if user:
                        if user.role == "admin":
                            # 管理员可以访问所有工作区
                            request.state.workarea_id = None
                        else:
                            # 检查用户是否属于请求的工作区
                            workarea_id = int(request.path_params.get("workarea_id", user.workarea_id))
                            if workarea_id != user.workarea_id:
                                return JSONResponse(
                                    status_code=403,
                                    content={"detail": "No permission to access this workarea"}
                                )
                            request.state.workarea_id = user.workarea_id
                finally:
                    db.close()
                    
            except JWTError as e:
                logger.error(f"JWT Error: {str(e)}")  # 添加错误日志
                return JSONResponse(
                    status_code=401,
                    content={"detail": "Invalid authentication token"}
                )
                
            response = await call_next(request)
            return response
            
        except Exception as e:
            logger.error(f"Middleware error: {str(e)}", exc_info=True)  # 添加详细错误日志
            return JSONResponse(
                status_code=500,
                content={"detail": f"Internal server error: {str(e)}"}
            ) 